'''
Created on January 11, 2010

This module contains all Disk Management related classes and functions 
'''

import subprocess
from os import popen
 

class DiskManagement(object) :
    '''handle disk management tasks like file systems , LVM , RAID , partitions and quotas.'''
    
    def __init__(self) :
        pass
    
    def listDevices(self):
        devices=popen("fdisk -l |grep dev | cut -d' ' -f1 |grep dev").read().strip().split("\n")
        return devices
    
    def listPartitionsInfo (self, devicePath = None) :
        """
        in this function if the hard disk device passed so it will get all partitions in the
            passed device , else then it will get all partitions in all devices in the system.
        """
        if (devicePath is not None):
            call = popen("parted -s "+devicePath+" print")
            devInfo = call.read()
            call.close()
            return devInfo
        
        else:
            call = popen("fdisk -l")
            devInfo = call.read()
            call.close()
            return devInfo
    
                     
    def deletePartition (self, devicePath = None ,partitionNumber = None) :
        """
        this function is very ########dangerous####### so take care before testing it
            First : if you passed the 'device path' and the "partition number" to the function your selected partition will be deleted.
            Second : if you passed the 'device path' only , all partitions related to selected device will be deleted so take care in your
            testing process.
        """
        if (partitionNumber is not None) & (devicePath is not None):
            if subprocess.call("parted -s "+devicePath+" rm "+partitionNumber , shell=True):
                return False
            else:
                return True
        
        elif (devicePath is not None):
            partitonsList = popen("parted -s "+devicePath+" print |awk '/^ / { print $1 }'").read()
            try:
                for loopPartition in partitonsList:
                    if subprocess.call("parted -s "+devicePath+" rm "+loopPartition , shell=True):
                        return False
                    else:
                        return True
            
            except Exception, e:
                return e
                     
    
    
    def formatPartition (self, partitionPath=None, formatType=None):
        """
        umount , format passed partition with passed format 
            Format Types : 'ext2' , 'ext3' , 'ext4' , 'jfs' , 'msdos' , 'ntfs' , 'vfat' and 'xfs '
            if you want to make swap partition just pass the 'partition path' and 'mkswap' to the function
        """
        if (formatType is not None) & (partitionPath is not None):
            if subprocess.call("umount "+partitionPath+" ; mkfs."+formatType+" "+partitionPath , shell=True):
                return False
            else:
                return True            
    
    
    def calculateSpace (self, devicePath=None) :
        """
        Get the size of the passed device.
        """
        if (devicePath is not None):
            if subprocess.call("parted -s "+devicePath+" print|awk '/^Disk/ { print $3 }'|sed 's/[Mm][Bb]//'" , shell=True):
                return False
            else:
                return True
    
    
    def resizePartition (self, devicePath=None , partitionNumber=None , startSize=None , endSize=None) :
        """
        in this function you will pass the device path as '/dev/sda' and pass the partition number the ,
            just the number of the partition in '/dev/sda(1)' we will passed '1' , and also passed the start size 
            in your hard disk as 1KB , 20MB , 30GB any thing and also the end size of the partition.
            - resize NTFS partitions not supported yet.
            - this function also exepted any error happened through resize process
        """
        if (devicePath is not None) & (partitionNumber is not None) & (startSize is not None) & (endSize is not None):
            try:
                if subprocess.call("parted -s "+devicePath+" resize "+partitionNumber+" "+startSize+"MB "+endSize+"MB" , shell=True):
                    return False
                else:
                    return True
            
            except Exception, e:
                return e
        
   
    def addNewPartiton (self, devicePath=None , partType=None , startSize=None , endSize=None) :
        """
        this function we create new partitions on known unlocated space (freespace)
            First : devicePath take the path of the for example '/dev/sdb'
            Second : take 3 types 
                1.'primary' and indentify the start size and the end size
                2.'extended' and also identify the start size and end size
                3.'logical' and we can't use this type until we already have 'extended' partition space to be able 
                    to add logical partitions to the extended space
        """
        if (devicePath is not None) & (partType is not None) & (startSize is not None) & (endSize is not None):
            try:
                if subprocess.call("parted -s "+devicePath+" mkpart "+partType+" "+startSize+"MB "+endSize+"MB" , shell=True):
                    return False
                else:
                    return True
            
            except Exception, e:
                return e
           
    
    def mountPartiton (self, partitonName , mountPoint , formatType =None , options = None , automate = None) :
        """
        in this functions we test the existency of the mount point if found then skip else create
            it , and then mount the passed device.
        """
        if (partitonName is not None) & (mountPoint is not None) & (options is not None) & (formatType is not None) & (automate is not None):
            if subprocess.call("[ -d "+mountPoint+" ] || mkdir "+mountPoint+" ; echo '"+partitonName+"    "+mountPoint+"    "+formatType+"    "+options+"    0 0 ' >> /etc/fstab ; mount -a",shell=True):
                return False
            else:
                return True

        elif (partitonName is not None) & (mountPoint is not None) & (formatType is not None) & (options is not None):
            if subprocess.call("[ -d "+mountPoint+" ] || mkdir "+mountPoint+" ; mount "+partitonName+" -t "+formatType+" -o "+options+" "+mountPoint ,shell=True):
                return False
            else:
                return True

        elif (partitonName is not None) & (mountPoint is not None) & (formatType is not None):
            if subprocess.call("[ -d "+mountPoint+" ] || mkdir "+mountPoint+" ; mount "+partitonName+" -t "+formatType+" "+mountPoint ,shell=True):
                return False
            else:
                return True

        elif (partitonName is not None) & (mountPoint is not None) & (options is not None):            
            if subprocess.call("[ -d "+mountPoint+" ] || mkdir "+mountPoint+" ; mount "+partitonName+" -t "+options+" "+mountPoint ,shell=True):
                return False
            else:
                return True

        elif (partitonName is not None) & (mountPoint is not None):
            if subprocess.call("[ -d "+mountPoint+" ] || mkdir "+mountPoint+" ; mount "+partitonName+" "+mountPoint ,shell=True):
                return False
            else:
                return True
            
        
    
    def unmountPartition (self, partitionPath=None, mountPoint=None) :
        """
        unmount passed device.
        """
        if (partitionPath is not None):
            if subprocess.call("umount "+partitionPath ,shell=True):
                return False
            else:
                return True
    
        elif (mountPoint is not None):
            if subprocess.call("umount "+mountPoint ,shell=True):
                return False
            else:
                return True
            
        elif (partitionPath is not None) & (mountPoint is not None):
            if subprocess.call("umount "+partitionPath+" "+mountPoint ,shell=True):
                return False
            else:
                return True

####################################################################################
################################### LVM ############################################
####################################################################################
    def prepareLVMPartition(self, devicePath=None , partNumber=None, state=None):
        """
        in this function we will set lvm flag to the partitions which we created previously
        we will path 3 variables:
            * the device path such as : /dev/sdb
            * the part number : 1,2,...
            * the part state : on/off
        """
        if (devicePath is not None) and (partNumber is not None) and (partNumber is not None):
            if subprocess.call("parted -s "+devicePath+" set "+partNumber+" lvm "+state+" ; partprobe ;", shell=True):
                return False
            else:
                return True
    
    
    def createPhysicalVolumes(self, LVMPart=None):
        """
        create physical volumes (physical extended) partitions , lvm partition will passed to the function
        """
        if (LVMPart is not None):
            if subprocess.call("pvcreate "+LVMPart , shell=True):
                return False
            else:
                return True
        
        
    def createInitialVolumGroup(self, VGName=None , PEPartition=None):
        """
        create initial VOLUM GROUP
        """
        if (VGName is not None) and (PEPartition is not None):
            if subprocess.call("vgcreate "+VGName+" "+PEPartition,shell=True):
                return False
            else:
                return True 
        
    def addPEtoVolumGroup(self,VGName=None,PEPartition=None):
        """
        add more physical extends to exist volume group 
        """
        if (VGName is not None) and (PEPartition is not None):
            if subprocess.call("vgextend "+VGName+" "+PEPartition,shell=True):
                return False
            else:
                return True
    
    def createLogicalVolum(self, size=None ,LVName=None, VGName=None):
        """
        create Logical Volumes
        """
        if (size is not None) and (LVName is not None) and (VGName is not None):
            if subprocess.call("lvcreate -L "+size+" -n "+LVName+" "+VGName,shell=True):
                return False
            else:
                return True
    
    
    def display(self, PVDisplay=None , VGDisplay=None , LVDisplay=None):
        if (PVDisplay is not None):
            pvcall = popen("pvdisplay")
            displaypv = pvcall.read()
            pvcall.close()
            return displaypv
        
        elif (VGDisplay is not None):
            vgcall = popen("vgdisplay")
            displayvg = vgcall.read()
            vgcall.close()
            return displayvg
            
        elif (LVDisplay is not None):
            lvcall = popen("lvdisplay")
            displaylv = lvcall.read()
            lvcall.close()
            return displaylv     
            
######################################################################################
#################################### RAID ############################################
######################################################################################
    def prepareRAIDPartition(self, devicePath=None , partNumber=None, state=None):
        """
        in this function we will set raid flag to the partitions which we created previously
        we will path 3 variables:
            * the device path such as : /dev/sdb
            * the part number : 1,2,...
            * the part state : on/off
        """
        if (devicePath is not None) and (partNumber is not None) and (partNumber is not None):
            if subprocess.call("parted -s "+devicePath+" set "+partNumber+" raid "+state+" ; partprobe ;", shell=True):
                return False
            else:
                return True
    
    
    def createRAIDPartitons(self,raidDevice=None,level=None,RAIDDevicesNumber=None,partitions=None):
        """
        in this function we will create raid devices , 4 variables will be passed
            * raidDevice : as /dev/md0 | /dev/md1 , ...etc
            * level : raid level 
            * RAIDDevicesNumber : number of partitions which will be added to raid device
            * partitions : the partitions and will be passed as string '/dev/sdb1 /dev/sdb2'.
        """
        if (raidDevice is not None) and (level is not None) and (RAIDDevicesNumber is not None) and (partitions is not None):
            if subprocess.call("mdadm --create "+raidDevice+" --level="+level+" --raid-devices="+RAIDDevicesNumber+" "+partitions,shell=True):
                return False
            else:
                return True
    
    def addPartitionsToRaidArray(self,raidDevice=None,partition=None):
        """
        add partiton to raidDevice
            * raidDevice : /dev/md0,...etc
            * partiton : /dev/sdb1 ,...etc
        """
        if (raidDevice is not None) and (partition is not None):
            if subprocess.call("mdadm "+raidDevice+" -a "+partition,shell=True):
                return False
            else:
                return True
    
    def removePartitonFromRaidArray(self,raidDevice=None,partition=None):
        """
        remove partiton to raidDevice
            * raidDevice : /dev/md0,...etc
            * partiton : /dev/sdb1 ,...etc
        """
        if (raidDevice is not None) and (partition is not None):
            if subprocess.call("mdadm "+raidDevice+" -f "+partition+" ; mdadm "+raidDevice+" -r "+partition,shell=True):
                return False
            else:
                return True
            
    def stopRaidArray(self, raidDevice=None):
        """
        stop raid array
            * raidDevice : /dev/md0,...etc
        """
        if (raidDevice is not None):
            if subprocess.call("mdadm --manage --stop "+raidDevice,shell=True):
                return False
            else:
                return True
    
    
    def getRAIDInfo(self,raidDevice=None):
        """
        get info of passed raid device
            * raidDevice : as /dev/md0 ,..etc
        """
        if (raidDevice is not None):
            call = popen("mdadm --detail "+raidDevice)
            devinfo = call.read()
            call.close()
            return devinfo
    
        
def main():
    # THE FOLLOWING CODE IS JUST FOR TESTING
    # CODER DISCRITION IS ADVISED :D
    x = DiskManagement()
    
    if x.chkbox_unmountPartition(partitionPath="/dev/sdb1"):
        print "#####Your device unmounted successfully#####"
    
    if x.mountPartiton(partitonName="/dev/sdb1", mountPoint="/media/FlashMemoryMP"):
        print "#####Your device mounted successfully#####"
    
    if x.chkbox_unmountPartition(mountPoint="/media/FlashMemoryMP"):
        print "#### MP Un mounted "
       
    if x.deletePartition(devicePath="/dev/sdb", partitionNumber="1"):
        print "#####partition deleted#####"
        
    if x.addNewPartiton(devicePath="/dev/sdb", partType="primary", startSize="1KB", endSize="1024MB"):
        print "#####primary partition with 1GB size added#####"
    
    if x.addNewPartiton(devicePath="/dev/sdb", partType="extended", startSize="1.1GB", endSize="3.5GB"):
        print "#####extended space added and waiting for add logical partitions to it#####"
    
    if x.addNewPartiton(devicePath="/dev/sdb", partType="logical", startSize="1.1GB", endSize="2GB"):
        print "#####first logical partition added to extended space#####"
        
    if x.addNewPartiton(devicePath="/dev/sdb", partType="logical", startSize="2.1GB", endSize="3.5GB"):
        print "#####second logical partition added to extended space#####"
    
    if x.deletePartition(devicePath="/dev/sdb", partitionNumber="1"):
        print "#####primary partition deleted successfully"

    if x.listPartitionsInfo(devicePath="/dev/sdb"):
        print "Got information of changes 1"
    
    if x.addNewPartiton(devicePath="/dev/sdb", partType="primary", startSize="1KB", endSize="900MB"):
        print "#####Re-add the primary partition########"
    
    if x.listPartitionsInfo(devicePath="/dev/sdb"):
        print "Got information of changes 2"
            
    if x.resizePartition(devicePath="/dev/sdb", partitionNumber="1", startSize="1KB", endSize="1024MB"):
        print "#####partition resized successfully#####"

    if x.listPartitionsInfo(devicePath="/dev/sdb"):
        print "Got information of changes 3"
                
    if x.formatPartition(formatType="ext3", partitionPath="/dev/sdb1"):
        print "#####/dev/sdb1 formated to ext3#####"
        
    if x.formatPartition(formatType="vfat", partitionPath="/dev/sdb5"):
        print "#####/dev/sdb5 formated to fat32#####"
    
    if x.formatPartition(mkswap="mkswap", partitionPath="/dev/sdb6"):
        print "#####swap partition added#####"
    
    if x.listPartitionsInfo(devicePath="/dev/sdb"):
        print "#####partition info listed again#####"
    
if __name__ == "__main__" :
    # Call the main function here to test this module
    #main()
    pass
    